{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mp.weixin.qq.com/s?__biz=MzkyMDMxNDkxOA==&mid=2247484354&idx=1&sn=50bc264fd92fe46221b03f548bfd9866&chksm=c195f68bf6e27f9dc47fda1cfafffe840684f40bfae4521a11210b54e0423e6bc71c1315b509&scene=178&cur_album_id=3447886516519747585#rd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from talib import MA_Type\n",
    "import talib\n",
    "from datetime import datetime, timedelta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hist_k_data(code,start_date,end_date,frequency='d')->pd.DataFrame:\n",
    "    \"\"\"\n",
    "    获取历史K线数据\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    import baostock as bs\n",
    "    bs.login()\n",
    "    rs = bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST\",start_date,end_date,frequency=frequency)\n",
    "    data_list = []\n",
    "    while (rs.error_code == '0') & rs.next():\n",
    "        # 获取一条记录，将记录合并在一起\n",
    "        data_list.append(rs.get_row_data())\n",
    "    result = pd.DataFrame(data_list, columns=rs.fields)\n",
    "    result.open = result.open.astype(float)\n",
    "    result.high = result.high.astype(float)\n",
    "    result.low = result.low.astype(float)\n",
    "    result.close = result.close.astype(float)\n",
    "    result.date= pd.to_datetime(result.date)\n",
    "    result.set_index('date',inplace=True)\n",
    "    bs.logout()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KDJ策略\n",
    " \n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    " \n",
    "# 假设df是包含股票历史数据的DataFrame，其中包含'Close'列\n",
    "# 计算KDJ指标\n",
    "def calculate_kdj(df, n=9, m=3):\n",
    "    df['H_n'] = df['Close'].rolling(window=n).max()\n",
    "    df['L_n'] = df['Close'].rolling(window=n).min()\n",
    "    df['RSV'] = (df['Close'] - df['L_n']) / (df['H_n'] - df['L_n']) * 100\n",
    "    df['K'] = df['RSV'].rolling(window=m).mean()\n",
    "    df['D'] = df['K'].rolling(window=m).mean()\n",
    "    df['J'] = 3 * df['K'] - 2 * df['D']\n",
    "    return df\n",
    " \n",
    "# 生成交易信号\n",
    "def generate_signals(df):\n",
    "    signals = pd.DataFrame(index=df.index)\n",
    "    signals['signal'] = 0.0\n",
    "    signals['signal'][n:] = np.where(df['K'][n:] > df['D'][n:], 1.0, 0.0) # 金叉\n",
    "    signals['positions'] = signals['signal'].diff()\n",
    "    return signals\n",
    " \n",
    "# 回测策略\n",
    "def backtest_strategy(df, signals):\n",
    "    df['positions'] = signals['positions']\n",
    "    df['strategy'] = df['positions'] * df['Close']\n",
    "    df['returns'] = df['strategy'].pct_change()\n",
    "    df['cumulative_returns'] = (1 + df['returns']).cumprod()\n",
    "    return df\n",
    " \n",
    "# 假设df是包含股票历史数据的DataFrame\n",
    "df = pd.read_csv('stock_data.csv')  # 读取股票数据\n",
    "df = df.set_index('Date')  # 设置日期为索引\n",
    "df = calculate_kdj(df)  # 计算KDJ指标\n",
    "signals = generate_signals(df)  # 生成交易信号\n",
    "df = backtest_strategy(df, signals)  # 回测策略\n",
    " \n",
    "# 绘制策略的累计收益\n",
    "plt.figure(figsize=(14, 7))\n",
    "plt.plot(df['cumulative_returns'], label='KDJ Strategy')\n",
    "plt.title('KDJ Strategy Cumulative Returns')\n",
    "plt.xlabel('Date')\n",
    "plt.ylabel('Cumulative Returns')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 假设df是包含股票历史数据的DataFrame，其中包含'Close'列\n",
    "# 计算KDJ指标\n",
    "def calculate_kdj(df, n=9, m=3):\n",
    "    df['H_n'] = df['close'].rolling(window=n).max()\n",
    "    df['L_n'] = df['close'].rolling(window=n).min()\n",
    "    df['RSV'] = (df['close'] - df['L_n']) / (df['H_n'] - df['L_n']) * 100\n",
    "    df['K'] = df['RSV'].rolling(window=m).mean()\n",
    "    df['D'] = df['K'].rolling(window=m).mean()\n",
    "    df['J'] = 3 * df['K'] - 2 * df['D']\n",
    "    return df\n",
    "\n",
    "# 生成交易信号\n",
    "def generate_signals(df):\n",
    "    signals = pd.DataFrame(index=df.index)\n",
    "    signals['signal'] = 0.0\n",
    "    signals['signal'][n:] = np.where(df['K'][n:] > df['D'][n:], 1.0, 0.0) # 金叉\n",
    "    signals['positions'] = signals['signal'].diff()\n",
    "    return signals\n",
    "\n",
    "# 回测策略\n",
    "def backtest_strategy(df, signals):\n",
    "    df['positions'] = signals['positions']\n",
    "    df['strategy'] = df['positions'] * df['close']\n",
    "    df['returns'] = df['strategy'].pct_change()\n",
    "    df['cumulative_returns'] = (1 + df['returns']).cumprod()\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义MACD策略参数\n",
    "symbol = 'sz.000001'\n",
    "\n",
    "# 获取K线数据\n",
    "end_time = '2024-07-02'\n",
    "start_time = '2005-01-01'\n",
    "df = get_hist_k_data(symbol, start_time, end_time)\n",
    "\n",
    "df = calculate_kdj(df)  # 计算KDJ指标\n",
    "signals = generate_signals(df)  # 生成交易信号\n",
    "df = backtest_strategy(df, signals)  # 回测策略"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 绘制策略的累计收益\n",
    "plt.figure(figsize=(14, 7))\n",
    "plt.plot(df['cumulative_returns'], label='KDJ Strategy')\n",
    "plt.title('KDJ Strategy Cumulative Returns')\n",
    "plt.xlabel('Date')\n",
    "plt.ylabel('Cumulative Returns')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
